#!/bin/bash

# Set default values
MODEL_NAME="unsloth/Meta-Llama-3.1-8B-Instruct"
LANG="rkt"
ROOT_DATASET="humaneval"
TEMPERATURE="0.2"
NUM_GPUS="2"
BATCH_SIZE="20"
COMPLETION_LIMIT="20"
OUTPUT_DIR_PREFIX="multiple_results"
TEMPLATE_TYPE="cot"

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --name)
            MODEL_NAME="$2"
            shift 2
            ;;
        --lang)
            LANG="$2"
            shift 2
            ;;
        --root-dataset)
            ROOT_DATASET="$2"
            shift 2
            ;;
        --temperature)
            TEMPERATURE="$2"
            shift 2
            ;;
        --num-gpus)
            NUM_GPUS="$2"
            shift 2
            ;;
        --template-type)
            TEMPLATE_TYPE="$2"
            shift 2
            ;;
        --tag)
            TAG="$2"
            shift 2
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Check if required arguments are provided
if [ -z "$LANG" ] || [ -z "$ROOT_DATASET" ] || [ -z "$TEMPERATURE" ]; then
    echo "Error: --lang, --root-dataset, and --temperature are required arguments."
    exit 1
fi

# template name = `template_type_lang`
TEMPLATE_NAME="${TEMPLATE_TYPE}_${LANG}"

# Set output directory
OUTPUT_DIR="$OUTPUT_DIR_PREFIX/$ROOT_DATASET/${MODEL_NAME//\//_}/${LANG}_${TEMPLATE_NAME}"

# prompt dataset path
PROMPT_DATASET_PATH="/home/sjw/ljb/lr_rag/multipl_e/prompts/${ROOT_DATASET}-${LANG}-reworded.jsonl"

# Create the output directory if it doesn't exist
mkdir -p "$OUTPUT_DIR"

# Run the Python script
# Check if the output directory already contains results
if [ -n "$(ls -A "$OUTPUT_DIR")" ]; then
    echo "Output directory $OUTPUT_DIR already contains results. Skipping inference step."
else
    echo "Running inference..."
    python inference_cot.py \
        --name "$MODEL_NAME" \
        --lang "$LANG" \
        --root-dataset "$ROOT_DATASET" \
        --temperature "$TEMPERATURE" \
        --output-dir "$OUTPUT_DIR" \
        --batch-size "$BATCH_SIZE" \
        --completion-limit "$COMPLETION_LIMIT" \
        --num_gpus "$NUM_GPUS" \
        --template-name "$TEMPLATE_NAME" \
        --use-local \
        --dataset "$PROMPT_DATASET_PATH"
fi

# run
python multipl_e/evaluation/src/main.py --dir "$OUTPUT_DIR" --output-dir "$OUTPUT_DIR" --full-function

# metrics
python multipl_e/pass_k.py "$OUTPUT_DIR"